import pandas as pd
import numpy as np
from enum import Enum, auto
from PIL import Image as PILImage

from ai_quickstart.openai_translator.ai_translator.utils import LOG


class ContentType(Enum):
    TEXT = auto()
    TABLE = auto()
    IMAGE = auto()


class Content:
    # 初始化内容类型、原文、翻译内容和状态
    def __init__(self, content_type, original, translation=None):
        self.content_type = content_type
        self.original = original
        self.translation = translation
        self.status = False

    # 设置翻译内容及状态
    def set_translation(self, translation, status):
        if not self.check_translation_type(translation):
            raise ValueError(f"Ivalid translation type. Expected {self.content_type}, but got {type(translation)}")
        self.translation = translation
        self.status = status

    # 检查翻译类型是否正确
    def check_translation_type(self, translation):
        if self.content_type == ContentType.TEXT and isinstance(translation, str):
            return True
        elif self.content_type == ContentType.TABLE and isinstance(translation, list):
            return True
        elif self.content_type == ContentType.IMAGE and isinstance(translation, PILImage.Image):
            return True
        return False


# 继承Content类，实现表格内容
class TableContent(Content):
    def __init__(self, data, translation=None):
        df = pd.DataFrame(data)

        # 判断数据是否为表格数据
        if len(data) != len(df) or len(data[0]) != len(df.columns):
            raise ValueError("The number of rows and columns in the extracted table data and dataframe do not match.")

        super().__init__(ContentType.TABLE, df)

    # 设置表格翻译内容及状态
    def set_translation(self, translation, status):
        try:
            # 判断翻译数据是否为文本数据，不是则抛出异常
            if not isinstance(translation, str):
                raise ValueError("Invalid translation type. Expected str, but got {type(translation)}")
            LOG.debug(translation)
            ### 使用mini-4o，返回的直接就是markdown的列表，可以直接使用
            # 将翻译数据按行分割，再按空格分割，形成二维列表
            # table_data = [row.strip().split() for row in translation.strip().split('\n')]
            # LOG.debug(table_data)
            #
            # # 判断翻译数据是否为DataFrame数据
            # translation_df = pd.DataFrame(table_data[1:], columns=table_data[0])
            # LOG.debug(translation_df)

            self.translation = translation
            self.status = status
        except Exception as e:
            LOG.error(f"Error in setting table translation: {e}")
            self.translation = None
            self.status = False

    # 重写__str__方法，返回原文的字符串形式，不包括表头和索引
    def __str__(self):
        return self.original.to_string(header=False, index=False)

    # 重写iter_items方法，返回行索引、列索引和单元格内容
    def iter_items(self, translated=False):
        # 根据翻译状态选择原文或翻译内容
        target_df = self.translation if translated else self.original
        # 遍历表格数据，返回行索引、列索引和单元格内容
        for row_idx, row in target_df.iterrows():
            for col_idx, item in enumerate(row):
                yield (row_idx, col_idx, item)

    # 更新表格内容
    def update_item(self, row_idx, col_idx, new_value, translated=False):
        # 根据翻译状态选择原文或翻译内容
        target_df = self.translation if translated else self.original
        # 更新表格内容
        target_df.at[row_idx, col_idx] = new_value

    # 获取原文表格的字符串形式，不包括表头和索引
    def get_original_as_str(self):
        return self.original.to_string(header=False, index=False)
